import random
import math
import matplotlib.pyplot as plt

def assign_cluster(x, centroids):
    """
    将样本 x 分配到最近的簇中心
    """
    min_dist = float('inf')
    idx = 0
    for i, c in enumerate(centroids):
        dist = math.dist(x, c)
        if dist < min_dist:
            min_dist = dist
            idx = i
    return idx


def Kmeans(data, k, epsilon=1e-4, iteration=100):
    """
    手动实现 K-means 聚类（不依赖任何工具包）
    """
    # 初始化 k 个随机中心
    centroids = random.sample(data, k)

    for _ in range(iteration):
        # 1.分配到最近簇
        clusters = [[] for _ in range(k)]
        for x in data:
            idx = assign_cluster(x, centroids)
            clusters[idx].append(x)

        # 2.计算新中心
        new_centroids = []
        for cluster in clusters:
            if len(cluster) == 0:
                new_centroids.append(random.choice(data))
            else:
                x_mean = sum(p[0] for p in cluster) / len(cluster)
                y_mean = sum(p[1] for p in cluster) / len(cluster)
                new_centroids.append((x_mean, y_mean))

        # 3.检查收敛
        shift = sum(math.dist(centroids[i], new_centroids[i]) for i in range(k))
        centroids = new_centroids
        if shift < epsilon:
            break

    return clusters, centroids

random.seed(0)

# 每类 50 个点
data = []
centers = [(2, 2), (7, 3), (4, 8)]
radius = 2  

for cx, cy in centers:
    for _ in range(50):    # 每类 50 个点
        x = random.uniform(cx - radius, cx + radius)
        y = random.uniform(cy - radius, cy + radius)
        data.append((x, y))


clusters, centroids = Kmeans(data, k=3)

print("最终中心点：")
for i, c in enumerate(centroids):
    print(f"Cluster {i}: ({c[0]:.3f}, {c[1]:.3f})")

plt.figure(figsize=(7,7))

# 画每一类
for cluster in clusters:
    xs = [p[0] for p in cluster]
    ys = [p[1] for p in cluster]
    plt.scatter(xs, ys, s=30)

# 画中心
cx = [c[0] for c in centroids]
cy = [c[1] for c in centroids]
plt.scatter(cx, cy, marker='X', s=200, linewidths=2)

plt.title("K-means Clustering Result")
plt.show()
